// Copyright 2016 The Fuchsia Authors
//
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT
#include "vm/vm_page_list.h"

#include <err.h>
#include <inttypes.h>
#include <trace.h>
#include <zircon/types.h>

#include <fbl/alloc_checker.h>
#include <ktl/move.h>
#include <vm/pmm.h>
#include <vm/vm.h>
#include <vm/vm_object_paged.h>

#include "vm_priv.h"

#define LOCAL_TRACE MAX(VM_GLOBAL_TRACE, 0)

namespace {

inline uint64_t offset_to_node_offset(uint64_t offset, uint64_t skew) {
  return ROUNDDOWN(offset + skew, PAGE_SIZE * VmPageListNode::kPageFanOut);
}

inline uint64_t offset_to_node_index(uint64_t offset, uint64_t skew) {
  return ((offset + skew) >> PAGE_SIZE_SHIFT) % VmPageListNode::kPageFanOut;
}

inline void move_vm_page_list_node(VmPageListNode* dest, VmPageListNode* src) {
  // Called by move ctor/assignment. Move assignment clears the dest node first.
  ASSERT(dest->IsEmpty());

  for (unsigned i = 0; i < VmPageListNode::kPageFanOut; i++) {
    vm_page* page = src->RemovePage(i);
    if (page) {
      dest->AddPage(page, i);
    }
  }
}

}  // namespace

VmPageListNode::VmPageListNode(uint64_t offset) : obj_offset_(offset) {
  LTRACEF("%p offset %#" PRIx64 "\n", this, obj_offset_);
}

VmPageListNode::~VmPageListNode() {
  LTRACEF("%p offset %#" PRIx64 "\n", this, obj_offset_);
  canary_.Assert();

  for (__UNUSED auto p : pages_) {
    DEBUG_ASSERT(p == nullptr);
  }
}

vm_page* VmPageListNode::GetPage(size_t index) const {
  canary_.Assert();
  DEBUG_ASSERT(index < kPageFanOut);
  return pages_[index];
}

vm_page* VmPageListNode::RemovePage(size_t index) {
  canary_.Assert();
  DEBUG_ASSERT(index < kPageFanOut);

  auto p = pages_[index];
  if (!p) {
    return nullptr;
  }

  pages_[index] = nullptr;

  return p;
}

zx_status_t VmPageListNode::AddPage(vm_page* p, size_t index) {
  canary_.Assert();
  DEBUG_ASSERT(index < kPageFanOut);
  if (pages_[index]) {
    return ZX_ERR_ALREADY_EXISTS;
  }
  pages_[index] = p;
  return ZX_OK;
}

VmPageList::VmPageList() { LTRACEF("%p\n", this); }

VmPageList::VmPageList(VmPageList&& other) : list_(std::move(other.list_)) {
  LTRACEF("%p\n", this);
  list_skew_ = other.list_skew_;
}

VmPageList::~VmPageList() {
  LTRACEF("%p\n", this);
  DEBUG_ASSERT(list_.is_empty());
}

VmPageList& VmPageList::operator=(VmPageList&& other) {
  list_ = std::move(other.list_);
  list_skew_ = other.list_skew_;
  return *this;
}

zx_status_t VmPageList::AddPage(vm_page* p, uint64_t offset) {
  uint64_t node_offset = offset_to_node_offset(offset, list_skew_);
  size_t index = offset_to_node_index(offset, list_skew_);

  if (node_offset >= VmObjectPaged::MAX_SIZE) {
    return ZX_ERR_OUT_OF_RANGE;
  }

  LTRACEF_LEVEL(2, "%p page %p, offset %#" PRIx64 " node_offset %#" PRIx64 " index %zu\n", this, p,
                offset, node_offset, index);

  // lookup the tree node that holds this page
  auto pln = list_.find(node_offset);
  if (!pln.IsValid()) {
    fbl::AllocChecker ac;
    ktl::unique_ptr<VmPageListNode> pl =
        ktl::unique_ptr<VmPageListNode>(new (&ac) VmPageListNode(node_offset));
    if (!ac.check()) {
      return ZX_ERR_NO_MEMORY;
    }

    LTRACEF("allocating new inner node %p\n", pl.get());
    __UNUSED auto status = pl->AddPage(p, index);
    DEBUG_ASSERT(status == ZX_OK);

    list_.insert(ktl::move(pl));
    return ZX_OK;
  } else {
    return pln->AddPage(p, index);
  }
}

vm_page* VmPageList::GetPage(uint64_t offset) const {
  uint64_t node_offset = offset_to_node_offset(offset, list_skew_);
  size_t index = offset_to_node_index(offset, list_skew_);

  LTRACEF_LEVEL(2, "%p offset %#" PRIx64 " node_offset %#" PRIx64 " index %zu\n", this, offset,
                node_offset, index);

  // lookup the tree node that holds this page
  auto pln = list_.find(node_offset);
  if (!pln.IsValid()) {
    return nullptr;
  }

  return pln->GetPage(index);
}

bool VmPageList::RemovePage(uint64_t offset, vm_page_t** page_out) {
  DEBUG_ASSERT(page_out);

  uint64_t node_offset = offset_to_node_offset(offset, list_skew_);
  size_t index = offset_to_node_index(offset, list_skew_);

  LTRACEF_LEVEL(2, "%p offset %#" PRIx64 " node_offset %#" PRIx64 " index %zu\n", this, offset,
                node_offset, index);

  // lookup the tree node that holds this page
  auto pln = list_.find(node_offset);
  if (!pln.IsValid()) {
    return false;
  }

  // free this page
  auto page = pln->RemovePage(index);
  if (page) {
    // if it was the last page in the node, remove the node from the tree
    if (pln->IsEmpty()) {
      LTRACEF_LEVEL(2, "%p freeing the list node\n", this);
      list_.erase(*pln);
    }

    *page_out = page;
    return true;
  } else {
    return false;
  }
}

void VmPageList::RemovePages(uint64_t start_offset, uint64_t end_offset,
                             list_node_t* removed_pages) {
  RemovePages([](vm_page_t*& p, uint64_t offset) -> bool { return true; }, start_offset, end_offset,
              removed_pages);
}

size_t VmPageList::RemoveAllPages(list_node_t* removed_pages) {
  LTRACEF("%p\n", this);

  DEBUG_ASSERT(removed_pages);

  size_t count = 0;

  // per page get a reference to the page pointer inside the page list node
  auto per_page_func = [&](vm_page*& p, uint64_t offset) {
    // add the page to our list and null out the inner node
    list_add_tail(removed_pages, &p->queue_node);
    p = nullptr;
    count++;
    return ZX_ERR_NEXT;
  };

  // walk the tree in order, freeing all the pages on every node
  ForEveryPage(per_page_func);

  // empty the tree
  list_.clear();

  return count;
}

bool VmPageList::IsEmpty() { return list_.is_empty(); }

void VmPageList::MergeFrom(VmPageList& other, const uint64_t offset, const uint64_t end_offset,
                           fbl::Function<void(vm_page*, uint64_t offset)> release_fn,
                           fbl::Function<bool(vm_page*, uint64_t offset)> migrate_fn,
                           list_node_t* free_list) {
  constexpr uint64_t kNodeSize = PAGE_SIZE * VmPageListNode::kPageFanOut;
  // The skewed |offset| in |other| must be equal to 0 skewed in |this|. This allows
  // nodes to moved directly between the lists, without having to worry about allocations.
  DEBUG_ASSERT((other.list_skew_ + offset) % kNodeSize == list_skew_);

  auto release_fn_wrapper = [&release_fn](vm_page_t*& p, uint64_t offset) -> bool {
    release_fn(p, offset);
    return true;
  };

  // Free pages outside of [|offset|, |end_offset|) so that the later code
  // doesn't have to worry about dealing with this.
  if (offset) {
    other.RemovePages(release_fn_wrapper, 0, offset, free_list);
  }
  other.RemovePages(release_fn_wrapper, end_offset, MAX_SIZE, free_list);

  // Calculate how much we need to shift nodes so that the node in |other| which contains
  // |offset| gets mapped to offset 0 in |this|.
  const uint64_t node_shift = offset_to_node_offset(offset, other.list_skew_);

  auto other_iter = other.list_.lower_bound(node_shift);
  while (other_iter.IsValid()) {
    uint64_t other_offset = other_iter->GetKey();
    // Any such nodes should have already been freed.
    DEBUG_ASSERT(other_offset < (end_offset + other.list_skew_));

    auto cur = other_iter++;
    auto other_node = other.list_.erase(cur);
    other_node->set_offset(other_offset - node_shift);

    auto target = list_.find(other_offset - node_shift);
    if (target.IsValid()) {
      // If there'a already a node at the desired location, then merge the two nodes.
      for (unsigned i = 0; i < VmPageListNode::kPageFanOut; i++) {
        auto page = other_node->RemovePage(i);
        if (!page) {
          continue;
        }

        zx_status_t status = target->AddPage(page, i);
        uint64_t src_offset = other_offset - other.list_skew_ + i * PAGE_SIZE;
        if (status == ZX_OK) {
          if (!migrate_fn(page, src_offset)) {
            target->RemovePage(i);
            list_add_tail(free_list, &page->queue_node);
          }
        } else {
          release_fn(page, src_offset);
          list_add_tail(free_list, &page->queue_node);
        }
      }
    } else {
      // If there's no node at the desired location, then directly insert the node.
      list_.insert_or_find(ktl::move(other_node), &target);
      bool has_page = false;
      for (unsigned i = 0; i < VmPageListNode::kPageFanOut; i++) {
        auto page = target->GetPage(i);
        if (page) {
          if (migrate_fn(page, other_offset - other.list_skew_ + i * PAGE_SIZE)) {
            has_page = true;
          } else {
            target->RemovePage(i);
            list_add_tail(free_list, &page->queue_node);
          }
        }
      }
      if (!has_page) {
        list_.erase(target);
      }
    }
  }
}

void VmPageList::MergeOnto(VmPageList& other, list_node_t* free_list) {
  DEBUG_ASSERT(other.list_skew_ == list_skew_);

  auto iter = list_.begin();
  while (iter.IsValid()) {
    auto node = list_.erase(iter++);
    auto target = other.list_.find(node->GetKey());
    if (target.IsValid()) {
      // If there'a already a node at the desired location, then merge the two nodes.
      for (unsigned i = 0; i < VmPageListNode::kPageFanOut; i++) {
        auto page = node->RemovePage(i);
        if (!page) {
          continue;
        }

        auto old_page = target->RemovePage(i);
        if (old_page) {
          list_add_tail(free_list, &old_page->queue_node);
        }

        target->AddPage(page, i);
      }
    } else {
      other.list_.insert(ktl::move(node));
    }
  }
}

VmPageSpliceList VmPageList::TakePages(uint64_t offset, uint64_t length) {
  VmPageSpliceList res(offset, length);
  const uint64_t end = offset + length;

  // Taking pages from children isn't supported, so list_skew_ should be 0.
  DEBUG_ASSERT(list_skew_ == 0);

  // If we can't take the whole node at the start of the range,
  // the shove the pages into the splice list head_ node.
  while (offset_to_node_index(offset, 0) != 0 && offset < end) {
    vm_page_t* page;
    if (RemovePage(offset, &page)) {
      res.head_.AddPage(page, offset_to_node_index(offset, 0));
    }
    offset += PAGE_SIZE;
  }

  // As long as the current and end node offsets are different, we
  // can just move the whole node into the splice list.
  while (offset_to_node_offset(offset, 0) != offset_to_node_offset(end, 0)) {
    ktl::unique_ptr<VmPageListNode> node = list_.erase(offset_to_node_offset(offset, 0));
    if (node) {
      res.middle_.insert(ktl::move(node));
    }
    offset += (PAGE_SIZE * VmPageListNode::kPageFanOut);
  }

  // Move any remaining pages into the splice list tail_ node.
  while (offset < end) {
    vm_page_t* page;
    if (RemovePage(offset, &page)) {
      res.tail_.AddPage(page, offset_to_node_index(offset, 0));
    }
    offset += PAGE_SIZE;
  }

  return res;
}

VmPageSpliceList::VmPageSpliceList() : VmPageSpliceList(0, 0) {}

VmPageSpliceList::VmPageSpliceList(uint64_t offset, uint64_t length)
    : offset_(offset), length_(length) {}

VmPageSpliceList::VmPageSpliceList(VmPageSpliceList&& other)
    : offset_(other.offset_),
      length_(other.length_),
      pos_(other.pos_),
      middle_(ktl::move(other.middle_)) {
  move_vm_page_list_node(&head_, &other.head_);
  move_vm_page_list_node(&tail_, &other.tail_);

  other.offset_ = other.length_ = other.pos_ = 0;
}

VmPageSpliceList& VmPageSpliceList::operator=(VmPageSpliceList&& other) {
  FreeAllPages();

  offset_ = other.offset_;
  length_ = other.length_;
  pos_ = other.pos_;
  move_vm_page_list_node(&head_, &other.head_);
  move_vm_page_list_node(&tail_, &other.tail_);
  middle_ = ktl::move(other.middle_);

  other.offset_ = other.length_ = other.pos_ = 0;

  return *this;
}

VmPageSpliceList::~VmPageSpliceList() { FreeAllPages(); }

void VmPageSpliceList::FreeAllPages() {
  // Free any pages owned by the splice list.
  while (!IsDone()) {
    auto page = Pop();
    if (page) {
      pmm_free_page(page);
    }
  }
}

vm_page* VmPageSpliceList::Pop() {
  if (IsDone()) {
    DEBUG_ASSERT_MSG(false, "Popped from empty splice list");
    return nullptr;
  }

  const uint64_t cur_offset = offset_ + pos_;
  const auto cur_node_idx = offset_to_node_index(cur_offset, 0);
  const auto cur_node_offset = offset_to_node_offset(cur_offset, 0);

  vm_page* res;
  if (offset_to_node_index(offset_, 0) != 0 &&
      offset_to_node_offset(offset_, 0) == cur_node_offset) {
    // If the original offset means that pages were placed in head_
    // and the current offset points to the same node, look there.
    res = head_.RemovePage(cur_node_idx);
  } else if (cur_node_offset != offset_to_node_offset(offset_ + length_, 0)) {
    // If the current offset isn't pointing to the tail node,
    // look in the middle tree.
    auto middle_node = middle_.find(cur_node_offset);
    if (middle_node.IsValid()) {
      res = middle_node->RemovePage(cur_node_idx);
    } else {
      res = nullptr;
    }
  } else {
    // If none of the other cases, we're in the tail_.
    res = tail_.RemovePage(cur_node_idx);
  }

  pos_ += PAGE_SIZE;
  return res;
}
